import paramiko
import threading
import time
import re
import logging
from queue import Queue
import json

# 配置详细日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('environment_validation.log')
    ]
)

class EnvironmentValidator:
    def __init__(self, host, username, password=None, key_path=None):
        self.host = host
        self.username = username
        self.password = password
        self.key_path = key_path
        self.ssh = None
        self.results = {}
        self.task_queue = Queue()
        self.completed_tasks = set()
        self.lock = threading.Lock()
        self.dependency_map = {}
        self.active_tasks = set()
        
    def connect(self):
        """建立SSH连接"""
        try:
            self.ssh = paramiko.SSHClient()
            self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            if self.key_path:
                private_key = paramiko.RSAKey.from_private_key_file(self.key_path)
                self.ssh.connect(self.host, username=self.username, pkey=private_key)
            else:
                self.ssh.connect(self.host, username=self.username, password=self.password)
            
            logging.info(f"✅ 成功连接到 {self.host}")
            return True
        except Exception as e:
            logging.error(f"❌ 连接失败: {str(e)}")
            return False
        
    def add_task(self, command, name=None, dependencies=None, inputs=None, timeout=30):
        """
        添加验证任务
        :param command: 要执行的命令
        :param name: 任务名称（用于依赖关系）
        :param dependencies: 依赖的任务名称列表
        :param inputs: 交互式输入（字典：{提示文本: 输入内容}）
        :param timeout: 命令超时时间（秒）
        """
        task_id = name if name else command
        task = {
            'command': command,
            'task_id': task_id,
            'dependencies': dependencies or [],
            'inputs': inputs or {},
            'timeout': timeout
        }
        
        # 添加到依赖映射
        for dep in task['dependencies']:
            if dep not in self.dependency_map:
                self.dependency_map[dep] = []
            self.dependency_map[dep].append(task_id)
        
        self.task_queue.put(task)
        logging.info(f"📝 添加任务: {task_id} (依赖: {dependencies})")

    def _wait_for_prompt(self, channel, timeout=5):
        """等待命令提示符出现，确保命令已结束"""
        start_time = time.time()
        output = ""
        
        while time.time() - start_time < timeout:
            if channel.recv_ready():
                data = channel.recv(1024).decode('utf-8', 'ignore')
                output += data
                
                # 检查是否出现标准提示符
                if re.search(r'[\$#]\s*$', output.splitlines()[-1] if output else ""):
                    return True, output
            
            time.sleep(0.1)
        
        return False, output

    def _execute_command(self, task):
        """执行单个命令并处理交互"""
        task_id = task['task_id']
        command = task['command']
        inputs = task['inputs']
        timeout = task['timeout']
        
        logging.info(f"🚀 开始执行: {task_id}")
        
        try:
            # 创建独立通道
            chan = self.ssh.invoke_shell()
            chan.settimeout(timeout)
            
            # 设置终端类型
            chan.send("export TERM=xterm\n")
            time.sleep(0.5)  # 给终端初始化时间
            
            # 发送命令
            chan.send(command + "\n")
            
            # 初始读取，获取命令回显
            time.sleep(0.5)
            initial_output = ""
            if chan.recv_ready():
                initial_output = chan.recv(4096).decode('utf-8', 'ignore')
            
            # 处理交互式输入
            output = initial_output
            start_time = time.time()
            last_data_time = time.time()
            
            while True:
                # 检查超时
                if time.time() - start_time > timeout:
                    raise TimeoutError(f"命令执行超时 ({timeout}秒)")
                
                # 检查是否有数据可读
                if chan.recv_ready():
                    data = chan.recv(4096).decode('utf-8', 'ignore')
                    output += data
                    last_data_time = time.time()
                    
                    # 调试：打印接收到的数据
                    logging.debug(f"[{task_id}] 接收: {data.strip()}")
                    
                    # 检查是否需要输入
                    input_sent = False
                    for prompt, response in inputs.items():
                        if prompt in data:
                            logging.debug(f"[{task_id}] 检测到提示: '{prompt}'")
                            chan.send(response + "\n")
                            logging.debug(f"[{task_id}] 发送响应: '{response}'")
                            input_sent = True
                            break
                    
                    # 如果发送了输入，继续等待输出
                    if input_sent:
                        continue
                
                # 检查命令是否已完成
                if chan.exit_status_ready():
                    exit_status = chan.recv_exit_status()
                    break
                
                # 检查是否长时间没有新数据（命令可能已完成）
                if time.time() - last_data_time > 2.0:  # 2秒无新数据
                    # 检查是否有命令提示符出现
                    prompt_found, prompt_output = self._wait_for_prompt(chan)
                    output += prompt_output
                    
                    if prompt_found:
                        logging.debug(f"[{task_id}] 检测到命令提示符，命令可能已完成")
                        exit_status = 0  # 假设成功
                        break
                    else:
                        # 尝试获取退出状态
                        if chan.exit_status_ready():
                            exit_status = chan.recv_exit_status()
                            break
                        else:
                            # 最终超时
                            raise TimeoutError("命令无响应，可能已完成或挂起")
                
                time.sleep(0.1)
            
            # 存储结果
            with self.lock:
                self.results[task_id] = {
                    'status': exit_status,
                    'output': output,
                    'command': command
                }
                self.completed_tasks.add(task_id)
                self.active_tasks.discard(task_id)
                
            logging.info(f"✅ 完成: {task_id} (状态: {exit_status})")
            
            # 标记依赖此任务的后续任务为可执行
            if task_id in self.dependency_map:
                for dependent_task in self.dependency_map[task_id]:
                    logging.debug(f"🔓 任务 {dependent_task} 的依赖 {task_id} 已满足")
                    
        except Exception as e:
            logging.error(f"❌ {task_id} 执行失败: {str(e)}")
            with self.lock:
                self.results[task_id] = {
                    'status': -1,
                    'output': f"执行错误: {str(e)}",
                    'command': command
                }
                self.completed_tasks.add(task_id)
                self.active_tasks.discard(task_id)
        finally:
            if chan:
                try:
                    chan.close()
                except:
                    pass

    def _task_worker(self):
        """工作线程处理任务"""
        while True:
            task = self.task_queue.get()
            if task is None:
                break
                
            task_id = task['task_id']
            
            # 检查依赖是否满足
            dependencies_met = all(dep in self.completed_tasks for dep in task['dependencies'])
            
            if dependencies_met or not task['dependencies']:
                # 添加到活动任务集
                with self.lock:
                    if task_id not in self.active_tasks and task_id not in self.completed_tasks:
                        self.active_tasks.add(task_id)
                    else:
                        # 任务已在进行中或已完成
                        self.task_queue.task_done()
                        continue
                
                # 执行任务
                self._execute_command(task)
            else:
                # 依赖未满足，重新放回队列
                self.task_queue.put(task)
                time.sleep(0.5)  # 避免CPU忙等待
                
            self.task_queue.task_done()

    def run_validation(self, num_workers=4):
        """运行验证任务"""
        if not self.ssh:
            if not self.connect():
                logging.error("无法建立SSH连接，终止验证")
                return False
        
        logging.info(f"🔧 开始环境验证，使用 {num_workers} 个工作线程...")
        
        # 创建工作线程
        workers = []
        for i in range(num_workers):
            worker = threading.Thread(target=self._task_worker)
            worker.daemon = True
            worker.start()
            workers.append(worker)
        
        # 等待所有任务完成
        self.task_queue.join()
        
        # 停止工作线程
        for _ in range(num_workers):
            self.task_queue.put(None)
        
        for worker in workers:
            worker.join()
        
        logging.info("\n✅ 所有验证任务已完成")
        return True

    def generate_report(self, output_file=None):
        """生成验证报告"""
        report = {
            "host": self.host,
            "username": self.username,
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            "results": self.results
        }
        
        # 控制台输出
        print("\n" + "="*50)
        print("环境验证报告".center(50))
        print("="*50)
        
        for task_id, result in self.results.items():
            status = "✅ 成功" if result['status'] == 0 else "❌ 失败"
            print(f"\n[{status}] {task_id}")
            print(f"命令: {result['command']}")
            
            # 清理输出中的ANSI转义码
            clean_output = re.sub(r'\x1B\[[0-?]*[ -/]*[@-~]', '', result['output'])
            
            print("输出摘要:")
            print(clean_output[:300].strip() + ("..." if len(clean_output) > 300 else ""))
            print("-"*50)
        
        # 文件输出
        if output_file:
            with open(output_file, 'w') as f:
                json.dump(report, f, indent=2)
            print(f"\n报告已保存到: {output_file}")
        
        return report

    def close(self):
        """关闭连接"""
        if self.ssh:
            try:
                self.ssh.close()
                logging.info("🔒 SSH连接已关闭")
            except:
                pass

# 示例使用
if __name__ == "__main__":
    # 配置连接信息
    validator = EnvironmentValidator(
        host="172.24.114.244",  # 可以是远程主机IP
        username="cc_ubuntu",
        password="789235"  # 或使用密钥认证
        # key_path="/path/to/private_key"
    )
    
    # 添加验证任务
    validator.add_task("echo '系统信息：' && uname -a", name="system_info")
    validator.add_task("echo '内存使用：' && free -h", name="memory_check")
    validator.add_task("echo '磁盘空间：' && df -h", name="disk_check")
    validator.add_task("echo '网络接口：' && ip addr", name="network_check")
    

    # 运行验证
    if validator.run_validation(num_workers=3):
        # 生成报告
        report = validator.generate_report("environment_report.json")
    
    # 清理
    validator.close()